package org.apache.seatunnel;

import sun.misc.BASE64Decoder;
import sun.misc.BASE64Encoder;

import javax.crypto.Cipher;

import java.io.File;
import java.io.IOException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/** RSA tool class (generate/save key pair, encrypt, decrypt) */
public class RSAUtils {

    /** algorithm name */
    private static final String ALGORITHM = "RSA";

    /** key length */
    private static final int KEY_SIZE = 2048;

    /** Randomly generate a key pair (including public and private keys) */
    public static KeyPair generateKeyPair() throws Exception {
        // Get the key pair generator for the specified algorithm
        KeyPairGenerator gen = KeyPairGenerator.getInstance(ALGORITHM);
        // Initialize key pair generator (specify key length, use default secure random number
        // source)
        gen.initialize(KEY_SIZE);
        // Randomly generate a pair of keys (including public key and private key)
        return gen.generateKeyPair();
    }

    /** Encode the public key/private key and save it to the specified file in Base64 format */
    public static void saveKeyForEncodedBase64(Key key, File keyFile) throws IOException {
        // returns the key in its primary encoding format
        byte[] encBytes = key.getEncoded();
        // Convert to Base64 text
        String encBase64 = new BASE64Encoder().encode(encBytes);
        // save to file
        IOUtils.writeFile(encBase64, keyFile);
    }

    /** Create a public key object from the Base64 text of the public key */
    public static PublicKey getPublicKey(String pubKeyBase64) throws Exception {
        // Convert the Base64 text of the public key to encoded public key bytes
        byte[] encPubKey = new BASE64Decoder().decodeBuffer(pubKeyBase64);
        // Create an encoded public key specification
        X509EncodedKeySpec encPubKeySpec = new X509EncodedKeySpec(encPubKey);
        // Obtain the key factory of the specified algorithm,
        // and generate a public key object according to the encoded public key specification
        return KeyFactory.getInstance(ALGORITHM).generatePublic(encPubKeySpec);
    }

    /** Create a private key object from the Base64 text of the private key */
    public static PrivateKey getPrivateKey(String priKeyBase64) throws Exception {
        // Convert the Base64 text of the private key to encoded private key bytes
        byte[] encPriKey = new BASE64Decoder().decodeBuffer(priKeyBase64);
        // Create an encoded private key specification
        PKCS8EncodedKeySpec encPriKeySpec = new PKCS8EncodedKeySpec(encPriKey);
        // Obtain the key factory of the specified algorithm,
        // and generate a private key object according to the encoded private key specification
        return KeyFactory.getInstance(ALGORITHM).generatePrivate(encPriKeySpec);
    }

    /** public key encrypted data */
    public static byte[] encrypt(byte[] plainData, PublicKey pubKey) throws Exception {
        // Get the cipher for the specified algorithm
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        // Initialize the cipher (public key encryption model)
        cipher.init(Cipher.ENCRYPT_MODE, pubKey);
        // Encrypted data, returns the encrypted ciphertext
        return cipher.doFinal(plainData);
    }

    /** private key to decrypt data */
    public static byte[] decrypt(byte[] cipherData, PrivateKey priKey) throws Exception {
        // Get the cipher for the specified algorithm
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        // Initialize the cipher (private key decryption model)
        cipher.init(Cipher.DECRYPT_MODE, priKey);
        // Decrypt the data and return the decrypted plaintext
        return cipher.doFinal(cipherData);
    }
}
